GH-47376: [C++][Compute] Support selective execution for kernels - #47377
GH-47376: [C++][Compute] Support selective execution for kernels#47377zanmato1984 wants to merge 8 commits into
Conversation
|
Thanks for opening a pull request! If this is not a minor PR. Could you open an issue for this pull request on GitHub? https://github.com/apache/arrow/issues/new/choose Opening GitHub issues ahead of time contributes to the Openness of the Apache Arrow project. Then could you also rename the pull request title in the following format? or See also: |
zanmato1984
changed the title
|
|
|
|
) ### Rationale for this change In order to support special form (#47374), the kernels have to respect the selection vector. Currently none of the kernels does. And it's almost impossible for us to make all existing kernels to respect the selection vector at once (and we probably never will). Thus we need an incremental way to add selection-vector-aware kernels on demand, meanwhile accommodate legacy (selection-vector-non-aware) kernels to be executed "selection-vector-aware"-ly in a general manner - the idea is to first "gather" selected rows from the batch into a new batch, evaluate the expression on the new batch, then "scatter" the result rows into the positions where they belong in the original batch. This makes the `take` and `scatter` functions dependencies of the exec facilities, which is in compute core (libarrow). And `take` is already in compute core. Now we need to move `scatter`. I'm implementing the selective execution of kernels in #47377, including invoking `take` and `scatter` as explained above. And I have to write tests of that in `exec_test.cc` which is deliberately declared to be NOT depending on libarrow_compute. ### What changes are included in this PR? Move scatter compute function into compute core. ### Are these changes tested? Yes. Manually tested. ### Are there any user-facing changes? None. * GitHub Issue: #47375 Authored-by: Rossi Sun <zanmato1984@gmail.com> Signed-off-by: Rossi Sun <zanmato1984@gmail.com>
0cce8c3 to
99a3217
Compare
99098d2 to
c005bbb
Compare
c502acc to
469ead1
Compare
469ead1 to
9a1c49e
Compare
|
Attaching benchmark results. The benchmark employs a trivial kernel that does nothing but spins specified number of times to simulate CPU intensity, and the number of rows of the batch is 4k. Baseline: Regular kernel, no selection vector. click to expandSparse: Selective kernel, with a selection vector of selectivity from 0 to 100%. click to expandDense: Regular kernel enclosed by gather/scatter, with a selection vector of selectivity from 0 to 100%. click to expand |
|
Some interesting comparisons to note:
|
|
Hi @pitrou @bkietz @westonpace @felipecrv , I know this is a big one, but I do hope some of you can help to review this PR - this is the most critical prerequisite for the if_else special form. Appreciated! |
966d999 to
9072a34
Compare
|
Kindly ping @felipecrv @pitrou . Thanks. |
|
Please kindly help @felipecrv @pitrou . Thanks! |
|
Really really need you guys' help :) @pitrou @felipecrv |
|
Ping @pitrou @felipecrv . I really want to move this forward, and a subsequent feature enabling special forms is awaiting. Please help, thanks! |
|
Hey @zanmato1984, sorry for not responding earlier. Reviewing this is a non-trivial task, especially as I'm worried about the API and efficiency implications of this. I'm also worried about what I perceive as additional complexity in the execution layer, which is already overly complicated with lots of weird cases (I'm afraid bugs and hidden inefficiencies will hide there). I'm also skeptical that a selection vector that is purely a vector of selection indices is a good solution performance-wise. In many cases you might have long spans of contiguous selected values, but with this solution we lose any opportunity of efficient batch execution (not to mention potential SIMD vectorization). This might be especially painful with chunked inputs (does this PR work with chunked arrays at all?). This is also my experience following the take/filter implementation work that we have been doing heroic efforts to try and improve the indirect indexing story while IMHO we should really be talking in terms of contiguous spans, or more precisely a mixed representation (spans/individual indices). All in all, I'm afraid this is committing to a lot of choices that both add complexity while painting us in a corner performance-wise. What do you think? You probably have thought about all this already. (also for the record my current professional situation is that I'm available for smaller maintenance and PR reviews but it's difficult to commit non-client time for such large enhancements) |
|
@pitrou Thanks a lot for taking the time and for the candid feedback - totally understood on the time constraints. I agree the execution layer is already overly complicated. In this PR I tried to keep the added complexity moderate and localized: when no selection vector is present, the existing execution path is unchanged (aside from a couple of null checks in the iterator). The selection-aware path is only activated when an ExecBatch carries a On performance: the selection path is only intended to be exercised from the upcoming special-form work (#47374), as a narrow and semantically explicit entry point. The primary goal there is semantic correctness under vectorized execution (e.g. guarding against errors like division-by-zero in rows where the condition is false), and I think paying some overhead is acceptable for that. I posted benchmark results earlier (#47377 (comment)); in summary, sparse execution wins strongly at low selectivity, while the worst regressions (up to ~4x) are from the generic dense fallback when a kernel doesn’t provide Chunked inputs: yes, this PR works with chunked arrays. Both the chunked + selection cases are covered by unit tests ( Finally, I hear you on representation: an index-only SelectionVector loses contiguous-span information and may limit batch/SIMD opportunities. I haven’t fully designed a mixed representation yet, but I agree it may be the right long-term direction. If you have a preferred API shape (hybrid runs+indices, a generalized “selection” object, or a different exec signature), I’d really appreciate guidance - I’d rather adjust before we cement the API. |
|
First I forgot another problem: the PR currently uses 32-bit indices, but we really want 64-bit indices, right? (perhaps UInt64 to match what sort_indices outputs, though that's unnecessary).
Separately from this PR, can we think about ways to make it simpler? Perhaps there are internal execution "options" that aren't really useful.
Yes, but we would like it to be more generally useful for execution engines, right?
The regression might be much worse on chunked inputs?
I'm not sure what it should look like, and we can probably add some complexity piecewise if we agree the API remains experimental. Ideally I'd like something that can be used internally for take/filter as well. A conceptual sketch could look like: struct ContiguousSpan {
int64_t start_offset;
int64_t length;
};
struct FilteredSpan {
int64_t start_offset;
int64_t length;
/* followed by a filter bitmap with `length` bits */
};
struct DiscreteSpan {
int64_t length;
/* followed by `length` 64-bit indices */
};
using SelectionSpan = std::variant<ContiguousSpan, FilteredSpan, DiscreteSpan>;(but SelectionSpan would actually be encoded using some bit-twiddling and a selection vector would be a Buffer containing a number of SelectionSpans) That's of course quite a bit of work and DiscreteSpan might be the only implemented variant at the start. |
- Deduplicate ApplySelectionMask for ArrayData/ArraySpan\n- Fix OUTPUT_NOT_NULL full-span validity init\n- Add unit test for OUTPUT_NOT_NULL selective exec
- Centralize SelectedCount/empty-span checks\n- Reduce scattered selection if/else in ScalarExecutor
apache#47378) ### Rationale for this change In order to support special form (apache#47374), the kernels have to respect the selection vector. Currently none of the kernels does. And it's almost impossible for us to make all existing kernels to respect the selection vector at once (and we probably never will). Thus we need an incremental way to add selection-vector-aware kernels on demand, meanwhile accommodate legacy (selection-vector-non-aware) kernels to be executed "selection-vector-aware"-ly in a general manner - the idea is to first "gather" selected rows from the batch into a new batch, evaluate the expression on the new batch, then "scatter" the result rows into the positions where they belong in the original batch. This makes the `take` and `scatter` functions dependencies of the exec facilities, which is in compute core (libarrow). And `take` is already in compute core. Now we need to move `scatter`. I'm implementing the selective execution of kernels in apache#47377, including invoking `take` and `scatter` as explained above. And I have to write tests of that in `exec_test.cc` which is deliberately declared to be NOT depending on libarrow_compute. ### What changes are included in this PR? Move scatter compute function into compute core. ### Are these changes tested? Yes. Manually tested. ### Are there any user-facing changes? None. * GitHub Issue: apache#47375 Authored-by: Rossi Sun <zanmato1984@gmail.com> Signed-off-by: Rossi Sun <zanmato1984@gmail.com>
github-actions
Bot
added
awaiting change review
fefe2e6 to
956c4f8
Compare
There was a problem hiding this comment.
Pull request overview
Adds “selective execution” support to Arrow C++ compute scalar kernels by introducing a selection-span abstraction, an optional selective kernel exec entrypoint, and a dense gather/scatter fallback for legacy kernels. This is foundational work for special-form expression evaluation (e.g. short-circuiting / non-strict evaluation) described in the linked issues.
Changes:
- Introduces
SelectionSpanvariants and expandsSelectionVectorto support chunk-relative spans, dense gather, and scatter. - Extends scalar kernel APIs with optional
ArrayKernelSelectiveExec, and updatesScalarExecutorto propagate/mask selection. - Adds dense selection fallback wrapper (
exec_selection.cc) plus extensive unit tests and a new benchmark.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| cpp/src/arrow/meson.build | Adds compute/exec_selection.cc to Meson build sources. |
| cpp/src/arrow/compute/test_util_internal.h | Adds selection-vector test helpers. |
| cpp/src/arrow/compute/test_util_internal.cc | Implements JSON/step-based selection-vector helpers. |
| cpp/src/arrow/compute/kernel.h | Adds ArrayKernelSelectiveExec and stores it on ScalarKernel. |
| cpp/src/arrow/compute/function.h | Adds ScalarFunction::AddKernel overload accepting selective exec. |
| cpp/src/arrow/compute/function.cc | Wires new AddKernel overload to store selective exec in kernels. |
| cpp/src/arrow/compute/expression.cc | Threads ExecBatch.selection_vector through scalar expression execution (debug-validated). |
| cpp/src/arrow/compute/exec.h | Introduces SelectionSpan and expands/refactors SelectionVector API; updates ExecBatch to carry selection separately from length. |
| cpp/src/arrow/compute/exec.cc | Implements index-backed SelectionVector, selection-aware ExecSpanIterator, selection masking in ScalarExecutor, and wraps scalar executor with dense-selection fallback. |
| cpp/src/arrow/compute/exec_test.cc | Adds coverage for SelectionVector validation/spans and for selective vs dense fallback execution behavior. |
| cpp/src/arrow/compute/exec_selection.cc | Implements dense gather/exec/scatter fallback wrapper for kernels without selective exec. |
| cpp/src/arrow/compute/exec_internal.h | Adds selection iteration helpers and declares dense-selection wrapper factory. |
| cpp/src/arrow/compute/exec_benchmark.cc | Adds benchmark measuring selective execution overhead patterns. |
| cpp/src/arrow/compute/CMakeLists.txt | Registers the new compute benchmark target. |
| cpp/src/arrow/CMakeLists.txt | Adds compute/exec_selection.cc to core Arrow compute sources. |
Suppressed comments (3)
cpp/src/arrow/compute/exec.cc:1610
IndexSelectionVector::ValidateusesARROW_CHECK_NE(indices_, nullptr), which hard-aborts the process on invalid/malformed inputs (e.g. ArrayData missing the values buffer). SinceSelectionVector::MakeIndices(ArrayData)is a public API entry point, this should return aStatus::Invalidinstead of terminating.
This issue also appears on line 1674 of the same file.
cpp/src/arrow/compute/exec.cc:1678
IndexSelectionVector::GetSpanForChunkdoes pointer arithmetic onindices_unconditionally (indices_ + selection_position). Ifindices_is null (possible with malformed input ArrayData, or if an empty indices array omits the values buffer), this is undefined behavior even whenlength() == 0. Add an early return for empty selections (and/or a defensive check forindices_).
cpp/src/arrow/compute/exec.h:191SelectionVectoris now abstract and can only be created viaMakeIndices(...), which is a source-breaking change for downstream code that previously constructed aSelectionVectordirectly (e.g.std::make_shared<SelectionVector>(array)in older versions). If this header is part of the public C++ API, consider providing a compatibility shim (deprecated helper / concrete type) or calling out the API break explicitly (the PR description says “None” for user-facing changes).
/// An index to represent that a batch does not belong to an ordered stream
constexpr int64_t kUnsequencedIndex = -1;
/// \brief A unit of work for kernel execution. It contains a collection of
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| /// that takes a selection vector argument and performs the computation only on the | ||
| /// selected indices. When this specialized kernel is not provided we fallback to | ||
| /// logic that gathers all selected values into a dense array, call `exec` on it | ||
| /// and then scather the values on the output array. |
| std::shared_ptr<SelectionVector> MakeSelectionVectorTo(int64_t length) { | ||
| auto res = gen::Step<int32_t>()->Generate(length); | ||
| DCHECK_OK(res.status()); | ||
| auto arr = res.ValueUnsafe(); | ||
| return SelectionVector::MakeIndices(*arr); | ||
| } |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 12 out of 12 changed files in this pull request and generated 3 comments.
Suppressed comments (3)
cpp/src/arrow/compute/exec.cc:488
- ExecSpanIterator::Next() dereferences selection_span unconditionally when selection_vector_ is set. Since selection_span is optional (default nullptr) and some call sites iterate spans without requesting selection, this will segfault in release when ExecBatch has a selection vector.
if (selection_vector_) {
DCHECK_NE(selection_span, nullptr);
auto indices_begin = selection_vector_->indices() + selection_position_;
auto indices_end = selection_vector_->indices() + selection_vector_->length();
DCHECK_LE(indices_begin, indices_end);
cpp/src/arrow/compute/exec.cc:1485
- SelectionVector::Validate() uses ARROW_CHECK_NE(indices_, nullptr), which aborts the process instead of returning an error status. Since Validate is part of the API contract, it should fail gracefully and also reject duplicate indices (current check allows non-decreasing order).
Status SelectionVector::Validate(int64_t values_length) const {
if (data_ == nullptr) {
return Status::Invalid("SelectionVector not initialized");
}
ARROW_CHECK_NE(indices_, nullptr);
if (data_->type->id() != Type::INT32) {
cpp/src/arrow/compute/exec.cc:914
- ExecuteSelectiveDense() calls Scatter() without passing the current ExecContext, which can allocate from a different/default context (wrong memory pool / CPU options) than the rest of execution.
ARROW_ASSIGN_OR_RAISE(auto result,
Scatter(dense_result, *batch.selection_vector->data(),
ScatterOptions{/*max_index=*/batch.length - 1}));
return listener->OnResult(std::move(result));
| if (!have_all_scalars_ || promote_if_all_scalars_) { | ||
| if (selection_vector_) { | ||
| DCHECK_NE(selection_span, nullptr); | ||
| *selection_span = SelectionVectorSpan(selection_vector_->indices()); | ||
| } | ||
| } |
| selection_vector_ = batch.selection_vector.get(); | ||
| if (selection_vector_) { | ||
| selection_length_ = selection_vector_->length(); | ||
| } else { | ||
| selection_length_ = 0; | ||
| } |
| } else if (kernel_->null_handling == NullHandling::INTERSECTION) { | ||
| if (!elide_validity_bitmap_) { | ||
| PropagateNullsSpans(input, result_span); | ||
| } | ||
| } else if (kernel_->null_handling == NullHandling::OUTPUT_NOT_NULL) { |
Rationale for this change
In order to support special form (#47374), being able to "selective"-ly execute the kernel becomes a prerequisite. As mentioned in #47374, we need an incremental way to add selective kernels on demand, meanwhile accommodate arbitrary legacy kernels to be executed selectively in a general manner.
What changes are included in this PR?
ArrayKernelSelectiveExec(KernelContext*, const ExecSpan&, const SelectionVectorSpan&, ExecResult*)in the kernel. This is the entry for selectively executing the kernel on a batch with a given selection vector. The kernel author can provide a dedicated implementation for such kernel API so the kernel can be executed "sparse"-ly - only the rows indicated by the selection vector will be processed. Otherwise the selective execution will fall back to a general "dense" way - gather the selected rows into a new contiguous (dense) batch, execute the kernel using the non-selective exec API, then scatter the result back to the original row positions.ScalarExecutorwith dense execution ability.Are these changes tested?
Tested and benchmarked.
Are there any user-facing changes?
None.